import unittest

import tensorflow as tf

from model.tf.resnet import resnet18, resnet50

HEIGHT, WIDTH, DEPTH = 64, 64, 3


class ResNetTestCase(unittest.TestCase):

    def setUp(self) -> None:
        gpus = tf.config.experimental.list_physical_devices('GPU')
        if gpus:
            try:
                tf.config.experimental.set_virtual_device_configuration(
                    gpus[0],
                    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=512)])
            except RuntimeError as e:
                print(e)

    def test_resnet18(self):
        resnet_18 = resnet18()
        resnet_18(tf.keras.layers.Input([HEIGHT, WIDTH, DEPTH]))
        print(resnet_18.summary())
        self.assertEqual((40, 200), resnet_18.compute_output_shape([40, HEIGHT, WIDTH, DEPTH]))

    def test_resnet50(self):
        resnet_50 = resnet50()
        resnet_50(tf.keras.layers.Input([HEIGHT, WIDTH, DEPTH]))
        print(resnet_50.summary())
        self.assertEqual((40, 200), resnet_50.compute_output_shape([40, HEIGHT, WIDTH, DEPTH]))